[trainer, worker, fsdp, cfg, tests] feat: add multi-role distillation runtime - #546

What does this PR do?

PR 2 of RFC #519: distributed role execution, independent student/fake-score optimizers and schedulers, LoRA/EMA, and composite checkpoint/resume. Runtime specification: #535.

Stacked on distillation-acceleration / #545; review the incremental diff against that parent. This is not a runnable model recipe: a separate architecture adapter must supply conditioning, forwards, score conversion, and differentiable student computation.

Checklist Before Starting

• Checked issue comments and searched PRs referencing #535 and multi-role distillation work.
• Title follows [{modules}] {type}: {description}.

This is #519's runtime layer, not a second controller or OPD trainer. It preserves OPD (#293, #495) and reuses existing trainer/FSDP foundations, adding independent fake-score optimization and multi-role state management.

Test

Local results at 28e2943, based on PR 1 1ef230a, using Python 3.12.13 / PyTorch 2.13.0+cu129 in the verlomni environment:

CPU coverage includes dispatch/futures, typed configs, sample/element accumulation, gradient ownership, finite-step accounting, EMA/export, and checkpoint fingerprints, using fake workers where appropriate.

Limits: historical one-rank FSDP1/FSDP2 evidence in #535 was not rerun for this head. The denominator collective still needs real multi-rank validation; no current-head production-training or inference-parity claim is made. Remote CPU CI only targets main / v0.* bases, so this stacked PR's CPU results are local.

export PYTHONPATH=. OMP_NUM_THREADS=2 TORCH_COMPILE_DISABLE=1 TORCHINDUCTOR_DISABLE=1
CUDA_VISIBLE_DEVICES='' python -m pytest -c /dev/null \
  -o 'python_files=*_on_cpu.py' -p no:cacheprovider --no-cov --asyncio-mode=auto -q tests/
CUDA_VISIBLE_DEVICES='' python -m pytest -c /dev/null \
  -p no:cacheprovider --no-cov --asyncio-mode=auto -q \
  tests/workers/test_diffusion_distillation_lora_on_cpu.py \
  tests/workers/test_diffusers_fsdp_merged_lora_on_cpu.py
CUDA_VISIBLE_DEVICES='' python -m pytest -c /dev/null \
  -p no:cacheprovider --no-cov --asyncio-mode=auto -q tests/special_sanity/test_config_docs.py
bash scripts/generate_trainer_config.sh
BASE=1ef230a880ee67ce8a45fc501e44b39904360e6d
mapfile -d '' files < <(git diff --cached --name-only --diff-filter=ACMR -z "$BASE")
python -m pre_commit run --files "${files[@]}"

API and Usage Example

Configuration fragment for a separately registered architecture adapter, not a standalone launch recipe:

algorithm:
  trainer_type: distillation
  sample_source: offline

actor_rollout_ref:
  actor:
    strategy: fsdp2
    optim:
      lr: 0.0001
      lr_scheduler_type: constant
  model:
    algorithm: dmd2
    lora_rank: 32

distillation:
  enabled: false
  distribution_matching:
    recipe: dmd2
    profile: distribution_only
    role_storage: shared_base_adapters
    fake_update_ratio: 2
    student_micro_batch_size_per_gpu: 1
    fake_score_micro_batch_size_per_gpu: 1
    fake_score_optim:
      lr: 0.00002
      lr_scheduler_type: constant
    ema_decay: 0.999
    ema_start_step: 0
    export_role: student

offline bypasses external RL rollout; student computation stays inside training. distillation.enabled remains OPD-only. Exporting student leaves EMA in resumable state.

The opt-in DistributionMatchingModelAdapter.build_distribution_matching_computer(model_config, plan) returns a DistributionMatchingComputer with:

• compute_phase(request, batch, runtime) → DistillationPhaseComputation: one graph-bearing scalar loss for the requested role plus detached metrics;
• state_dict() / load_state_dict(state): checkpointable sampling/RNG state.

The computer owns model math; the runtime owns optimization. Existing DiffusionModelBase adapters gain no new mandatory abstract methods.

Design & Code Changes

main_diffusion / TaskRunner
  -> DistillationRayTrainer(BaseRayDiffusionTrainer)
     -> PR 1 DistillationTrainerController
        -> DiffusionDistillationWorkerGroup (executor facade)
           -> Ray worker / DistributionMatchingComputer
              -> DistillationRoleRuntime
                 -> one DistillationRoleGroupEngine per physical group
                    -> existing DiffusersFSDPEngine

Ownership and reuse

Paths are relative to verl_omni/. The worker reuses verl dispatch/profiling and EngineRegistry, not a new loading backend or PPO-shaped execution.

Roles vs. storage: shared_base_adapters uses one frozen base with student/fake-score/EMA adapters and an adapter-disabled teacher. colocated_independent uses separate physical groups as the full-module fallback/correctness baseline. Sharing is not an algorithm requirement. Both layouts are colocated; shared storage requires LoRA, and shared FSDP1 requires use_orig_params=true.

One optimizer owner per phase: accumulate micro-batches, reject inactive-role gradients, normalize/clip, then step. Scheduler and eligible EMA updates require a finite successful step. Student/fake micro-batches and scheduler horizons are independent; PPO batch fields are not repurposed.

With loss_normalizer=None, accumulation uses sample means. An explicit positive count accumulates numerator gradients and divides by the DP-averaged denominator before clipping. Invalid/mixed modes fail closed; this does not imply globally element-weighted aggregation for every metric.

Architecture boundary: during a student phase, no-grad teacher/fake scoring must preserve the pending student graph; the score branches are detached. Role contexts restore adapter and train/eval state. Phase RPCs are blocking; computers still own synchronization of collective-dependent forward counts and branches.

Checkpoint vs. export: completed-cycle checkpoints atomically publish role groups, optimizers/schedulers, EMA, computer/driver RNG, dataloader and controller state. Persisted models are saved once per physical group; frozen teacher-only groups reload their base. Canonical fingerprints reject plan/config/layout drift. Publication is atomic; failed updates/restores are not in-memory rollback transactions.

Export selects student or student_ema with its own PEFT config and restores the previous adapter. Upstream merged-default export remains intact; named adapters are not silently substituted. This does not implement CheckpointEngine transfers or validation replicas.

Not included: architecture-specific training/generation, adversarial multi-optimizer phases, causal/KV-cache execution, standalone score transport, inference orchestration, HDFS restore, or automatic checkpoint retention. Model integration begins with #543; later RFC stages supply the remaining capabilities.

Checklist Before Submitting

• Read the Contribute Guide.
• Applied pre-commit checks to the complete incremental PR file set.
• Documented ownership, configuration and checkpoint contracts in [RFC] Multi-role distillation data plane, LoRA EMA, and checkpointing #535, building on [RFC] Unified Distribution-Matching Distillation Runtime for DMD, DMD2, CausVid, and Self-Forcing #519/[RFC] Generic distillation trainer control plane, contracts, and pure DMD math #520.
• Added CPU regressions to existing L1 discovery and separate GPU tests outside L1; validation limits are stated above.
• Reconfirm human review of the latest code delta before merging.

AI assistance: Claude and OpenAI via pi assisted implementation, tests, and review. The human submitter previously confirmed line-by-line review and testing of the earlier changes; the latest delta still needs final review.